// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
// Author: Markus Gutschke

#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#include "base64.h"
#include "check.h"
#include "net.h"
#include "stringprintf.h"
#include "xmpp.h"

#define min(a,b) ({ typeof(a) _a=(a); typeof(b) _b=(b); _a < _b ? _a : _b; })


static struct XMLElement *XML_NONE = (struct XMLElement *)-1;


static void startStream(struct XMPP *xmpp) {
  static const char hello[] =
    "<stream:stream to=\"%s\" xml:lang=\"en\" version=\"1.0\" "
        "xmlns:stream=\"http://etherx.jabber.org/streams\" "
        "xmlns=\"jabber:client\">";
  char *domain = xmlEscape(xmpp->server);
  char *msg = stringPrintf(NULL, hello, domain);
  free(domain);
  size_t msg_len = strlen(msg);
  check(connWrite(xmpp->conn, msg, msg_len) == msg_len);
  free(msg);
  free(xmpp->id);
  xmpp->id = NULL;
}

static void startSASL(struct XMPP *xmpp) {
  // Server advertised OAuth2. Initiate SASL handshake.
  char *raw_credentials, *credentials;
  size_t jid_len        = strlen(xmpp->jid);
  size_t auth_len       = strlen(xmpp->oauth2->accessToken);
  size_t raw_len        = jid_len + auth_len + 2;
  size_t b64_len        = (raw_len + 2)/3*4;
  check(raw_credentials = malloc(raw_len + b64_len + 1));
  credentials           = raw_credentials + raw_len;

  // Assemble credentials
  *raw_credentials      = '\000';
  strcpy(raw_credentials + 1, xmpp->jid);
  memcpy(raw_credentials + jid_len + 2,
         xmpp->oauth2->accessToken, auth_len);

  // Convert to BASE64
  check(b64encode(credentials, raw_credentials, raw_len)==b64_len);
  
  // Assemble SASL XML stanza.
  char *sasl = stringPrintf(NULL,
    "<auth xmlns=\"urn:ietf:params:xml:ns:xmpp-sasl\" "
          "mechanism=\"X-OAUTH2\" auth:service=\"chromiumsync\" "
          "auth:allow-generated-jid=\"true\" "
          "auth:client-uses-full-bind-result=\"true\" "
          "xmlns:auth=\"http://www.google.com/talk/protocol/auth\">"
      "%s"
    "</auth>",
    credentials);
  free(raw_credentials);

  // Send it to the server to initiate SASL handshake.
  size_t sasl_len = strlen(sasl);
  check(connWrite(xmpp->conn, sasl, sasl_len) == sasl_len);
  free(sasl);
}

static void restartWithFreshAuthorization(struct XMPP *xmpp) {
  refreshOAuth2(xmpp->oauth2, NULL);
  check(xmpp->conn->fd >= 0 && xmpp->conn->ssl && xmpp->conn->handle);
  deleteConnection(xmpp->conn);
  xmpp->conn = NULL;
  xmpp->restartStream = 1;
}

static void bindPrinter(struct XMPP *xmpp) {
  char *desc = xmlEscape(xmpp->printerDesc);
  char *bind = stringPrintf(NULL,
    "<iq type=\"set\" id=\"0\">"
      "<bind xmlns=\"urn:ietf:params:xml:ns:xmpp-bind\">"
        "<resource>%s</resource>"
      "</bind>"
    "</iq>",
    desc);
  free(desc);
  size_t bind_len = strlen(bind);
  check(connWrite(xmpp->conn, bind, bind_len) == bind_len);
  free(bind);
}

static void startSession(struct XMPP *xmpp, const char *fulljid) {
  free(xmpp->fulljid);
  free(xmpp->jid);
  check(xmpp->fulljid = strdup(fulljid));
  size_t len = strcspn(xmpp->fulljid, "/");
  check(xmpp->jid = malloc(len + 1));
  memcpy(xmpp->jid, xmpp->fulljid, len);
  xmpp->jid[len] = '\000';
  static const char session[] =
    "<iq type=\"set\" id=\"1\">"
      "<session xmlns=\"urn:ietf:params:xml:ns:xmpp-session\"/>"
    "</iq>";
  check(connWrite(xmpp->conn, session, sizeof(session)-1)==sizeof(session)-1);
}

static void subscribePrinter(struct XMPP *xmpp) {
  char *barejid = xmlEscape(xmpp->jid);
  char *msg = stringPrintf(NULL,
    "<iq type=\"set\" to=\"%s\" id=\"2\">"
      "<subscribe xmlns=\"google:push\">"
        "<item channel=\"cloudprint.google.com\" "
              "from=\"cloudprint.google.com\"/>"
      "</subscribe>"
    "</iq>",
    barejid);
  free(barejid);
  size_t msg_len = strlen(msg);
  check(connWrite(xmpp->conn, msg, msg_len) == msg_len);
  free(msg);
}

static int handshake(struct XML *xml, struct XMLElement *elem, void *arg) {
  struct XMPP *xmpp = (struct XMPP *)arg;
  check(elem); // Stream was closed prematurely
  struct XMLElement *child;
  if (elem == XML_NONE) {
    // Start the XML stream. This happens several times during the handshake.
    startStream(xmpp);
    return 1;
  } else if (xmlElementMatches(elem, "http://etherx.jabber.org/streams",
                               "stream")) {
    // Remember the "id" returned in the <stream> element
    const char *id;
    check(id = getFromHashMap(&elem->attributes, "id"));
    check(!xmpp->id);
    check(xmpp->id = strdup(id));
    return 1;
  } else if (xmlElementMatches(elem, "http://etherx.jabber.org/streams",
                               "features")) {
    // Parse the <features> and handle TLS, SASL and session negotiations
    if (xmlElementHasChild(elem, NULL, "urn:ietf:params:xml:ns:xmpp-tls",
                           "starttls", NULL)) {
      // Negotiate TLS
      static const char starttls[] =
        "<starttls xmlns=\"urn:ietf:params:xml:ns:xmpp-tls\"/>";
      check(connWrite(xmpp->conn, starttls, sizeof(starttls)-1) ==
            sizeof(starttls)-1);
      return 1;
    } else if (xmlElementHasChild(elem, &child,
                                  "urn:ietf:params:xml:ns:xmpp-sasl",
                                  "mechanisms", NULL) &&
               xmlElementHasChild(child, NULL,
                                  "urn:ietf:params:xml:ns:xmpp-sasl",
                                  "mechanism", "X-OAUTH2")) {
      // Negotiate SASL
      check(connectionIsSecure(xmpp->conn));
      startSASL(xmpp);
      return 1;
    } else if (xmlElementHasChild(elem, NULL,
                                  "urn:ietf:params:xml:ns:xmpp-bind",
                                  "bind", NULL) &&
               xmlElementHasChild(elem, NULL,
                                  "urn:ietf:params:xml:ns:xmpp-session",
                                  "session", NULL)) {
      // Bind the printer
      check(connectionIsSecure(xmpp->conn));
      bindPrinter(xmpp);
      return 1;
    }
  } else if (xmlElementMatches(elem, "urn:ietf:params:xml:ns:xmpp-tls",
                               "proceed")) {
    // TLS handshake is ready to go ahead. Reset XML stream.
    xmpp->negotiateSSL  = 1;
    xmpp->restartStream = 1;
    return 0;
  } else if (xmlElementMatches(elem, "urn:ietf:params:xml:ns:xmpp-sasl",
                               "success")) {
    // SASL negotation completed successfully. Reset XML stream.
    check(connectionIsSecure(xmpp->conn));
    xmpp->restartStream = 1;
    return 0;
  } else if (xmlElementMatches(elem, "urn:ietf:params:xml:ns:xmpp-sasl",
                               "failure")) {
    // SASL negotiation failed. Restart connection with new auth token.
    restartWithFreshAuthorization(xmpp);
    return 0;
  } else if (xmlElementMatches(elem, "jabber:client", "iq") &&
             xmlAttributeMatches(elem, "type", "result")) {
    check(connectionIsSecure(xmpp->conn));
    if (xmlAttributeMatches(elem, "id", "0") &&
        xmlElementHasChild(elem, &child, "urn:ietf:params:xml:ns:xmpp-bind",
                           "bind", NULL) &&
        xmlElementHasChild(child, &child, "urn:ietf:params:xml:ns:xmpp-bind",
                           "jid", NULL) &&
        child->children_size == 1 && child->children[0]->type == X_CDATA) {
      startSession(xmpp, child->children[0]->cdata);
      return 1;
    } else if (xmlAttributeMatches(elem, "id", "1")) {
      subscribePrinter(xmpp);
      return 1;
    } else if (xmlAttributeMatches(elem, "id", "2") &&
               xmlAttributeMatches(elem, "to", xmpp->fulljid)) {
      xmpp->completed = 1;
      return 0;
    }
  }
  check(0);
}

static void connectToXMPPServer(struct XMPP *xmpp) {
  for (xmpp->completed = 0; !xmpp->completed; ) {
    // During TLS and SASL negotiation, the XML stream resets a couple of
    // times. We have to call resetXML() as we don't want to discard any
    // extra bytes that might already have been read and buffered in the
    // XMP object at this point.
    xmpp->restartStream = 0;
    resetXML(&xmpp->xml);
    if (!xmpp->conn) {
    fixed_hostname:;
      int fd;
      if ((fd = openSocket(xmpp->server, "xmpp-client")) < 0 &&
          !strcmp(xmpp->server, "cloudprint.googleusercontent.com")) {
        // Work-around for temporary host name during early roll-out of
        // the Google CloudPrint service.
        strcpy(xmpp->server, "gmail.com");
        goto fixed_hostname;
      }
      check(fd >= 0);
      xmpp->conn = newConnection(fd);
    }
    handshake(&xmpp->xml, XML_NONE, xmpp);
    while (!xmpp->completed) {
      char buf[256];
      ssize_t sz;
      if (connectionIsSecure(xmpp->conn)) {
        check((sz = connRead(xmpp->conn, buf, sizeof(buf))) > 0);
      } else {
        // While in plain text mode, we always peek into the stream first,
        // and then consume bytes as we are successfully parsing them. This
        // allows the callback to notify us that it is time to switch to TLS.
        // The switch has to be done when the file descriptor is at exactly
        // the correct position.
        check((sz = connPeek(xmpp->conn, buf, sizeof(buf))) > 0);
      }
      ssize_t consumed = parseXML(&xmpp->xml, buf, sz, handshake, xmpp);
      if (xmpp->conn && !connectionIsSecure(xmpp->conn)) {
        if (xmpp->negotiateSSL) {
          // We are switching to TLS. Only read exactly the bytes that are
          // plain text, then perform the switch
          sz = consumed;
        }
        for (ssize_t rc; sz > 0; sz -= rc) {
          check((rc = connRead(xmpp->conn, buf, min(sizeof(buf), sz)))>0);
        }
        if (xmpp->negotiateSSL) {
          xmpp->negotiateSSL = 0;
          upgradeConnection(xmpp->conn);
        }
      }
      if (xmpp->restartStream) {
        break;
      }
    }
  }
}

int handleXmppPacket(struct XMPP *xmpp, int (*fnc)(struct XML *xml,
                       struct XMLElement *elem, void *arg), void *arg) {
  ssize_t sz;
  for (char buf[256]; (sz = connRead(xmpp->conn, buf, sizeof(buf))) > 0; ) {
    parseXML(&xmpp->xml, buf, sz, fnc, arg);
    if (!connIsReadable(xmpp->conn, 0)) {
      break;
    }
  }
  return sz != 0;
}

int tickleXmpp(struct XMPP *xmpp) {
  if (connIsWriteable(xmpp->conn, 0)) {
    char *fulljid = xmlEscape(xmpp->fulljid);
    char *ping = stringPrintf(NULL,
                              "<iq from='%s' id='ping%u' type='get'>"
                              "<ping xmlns='urn:xmpp:ping'/></iq>",
                              fulljid, ++xmpp->ping);
    free(fulljid);
    if (connWrite(xmpp->conn, ping, strlen(ping)) < 0) {
      free(ping);
      return -1;
    }
    free(ping);
  }
  return 0;
}

struct XMPP *newXMPP(struct OAuth2 *oauth2, const char *jid,
                     const char *printerId, const char *printerDesc) {
  struct XMPP *xmpp;
  check(xmpp = malloc(sizeof(struct XMPP)));
  initXMPP(xmpp, oauth2, jid, printerId, printerDesc);
  return xmpp;
}

void initXMPP(struct XMPP *xmpp, struct OAuth2 *oauth2, const char *jid,
              const char *printerId, const char *printerDesc) {
  memset(xmpp, 0, sizeof(struct XMPP));

  // OAuth2 object
  xmpp->oauth2 = oauth2;

  // Printer description
  check(xmpp->printerDesc = strdup(printerDesc));

  // XMPP jid
  xmpp->fulljid = NULL;
  check(xmpp->jid = strdup(jid));

  // Extract domain name from jid
  char *ptr;
  check(ptr = strchr(xmpp->jid, '@'));
  check(xmpp->server = strdup(ptr + 1));
  if ((ptr = strchr(xmpp->server, '/')) != NULL) {
    *ptr = '\000';
  }

  // Printer id
  check(xmpp->printerId = strdup(printerId));

  // Connect to server and perform initial handshake
  xmpp->conn = NULL;
  initXML(&xmpp->xml);
  connectToXMPPServer(xmpp);
}

void destroyXMPP(struct XMPP *xmpp) {
  if (xmpp) {
    deleteOAuth2(xmpp->oauth2);
    free(xmpp->server);
    free(xmpp->jid);
    free(xmpp->fulljid);
    free(xmpp->printerId);
    free(xmpp->printerDesc);
    free(xmpp->id);
    deleteConnection(xmpp->conn);
    destroyXML(&xmpp->xml);
  }
}

void deleteXMPP(struct XMPP *xmpp) {
  destroyXMPP(xmpp);
  free(xmpp);
}
